# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

from enum import IntEnum
import math
from coin.base.book.book_builder_base import (BookBuilderBase, PriceQty)


class BookUpdateSide(IntEnum):
  BID = 0
  ASK = 1


def to_aggregated_price_bid(price, ticksize, price_decimal=2):
  return round(math.floor(price / ticksize) * ticksize, price_decimal)


def to_aggregated_price_ask(price, ticksize, price_decimal=2):
  return round(math.ceil(price / ticksize) * ticksize, price_decimal)


class AggregatedBookBuilder(BookBuilderBase):
  def __init__(self, original_book: BookBuilderBase, aggregated_ticksize):
    super().__init__(original_book.exchange, original_book.symbol)
    self._original_book = original_book
    self._aggregated_ticksize = aggregated_ticksize
    self._price_maps = [self._bid_price_map, self._ask_price_map]

  def clear(self):
    self._price_maps[0].clear()
    self._price_maps[1].clear()

  def generate_aggregated_book(self, update_side: BookUpdateSide, qty_decimal=8):
    self.request_timestamp = self._original_book.request_timestamp
    self.fetched_timestamp = self._original_book.fetched_timestamp
    self.exchange_timestamp = self._original_book.exchange_timestamp
    self._price_maps[update_side].clear()

    if update_side == BookUpdateSide.BID:
      if self._original_book.bid0() is None:
        return
      try:
        curr_aggregated_price = to_aggregated_price_bid(self._original_book.bid0().price,
                                                        self._aggregated_ticksize)
      except IndexError:
        return
      curr_aggregated_qty = 0
      for price, price_qty in self._original_book.iteritems_bid():
        qty = price_qty.qty
        aggregated_price = to_aggregated_price_bid(price, self._aggregated_ticksize)
        if aggregated_price == curr_aggregated_price:
          curr_aggregated_qty += qty
        else:
          curr_aggregated_qty = round(curr_aggregated_qty, qty_decimal)
          if curr_aggregated_qty > 0:
            self._price_maps[update_side][curr_aggregated_price] = \
                PriceQty(curr_aggregated_price, curr_aggregated_qty)
          curr_aggregated_price = aggregated_price
          curr_aggregated_qty = qty
      curr_aggregated_qty = round(curr_aggregated_qty, qty_decimal)
      if curr_aggregated_qty > 0:
        self._price_maps[update_side][curr_aggregated_price] = \
            PriceQty(curr_aggregated_price, curr_aggregated_qty)
    elif update_side == BookUpdateSide.ASK:
      if self._original_book.ask0() is None:
        return
      try:
        curr_aggregated_price = to_aggregated_price_ask(self._original_book.ask0().price,
                                                        self._aggregated_ticksize)
      except IndexError:
        return
      curr_aggregated_qty = 0
      for price, price_qty in self._original_book.iteritems_ask():
        qty = price_qty.qty
        aggregated_price = to_aggregated_price_ask(price, self._aggregated_ticksize)
        if aggregated_price == curr_aggregated_price:
          curr_aggregated_qty += qty
        else:
          curr_aggregated_qty = round(curr_aggregated_qty, qty_decimal)
          if curr_aggregated_qty > 0:
            self._price_maps[update_side][curr_aggregated_price] = \
                PriceQty(curr_aggregated_price, curr_aggregated_qty)
          curr_aggregated_price = aggregated_price
          curr_aggregated_qty = qty
      curr_aggregated_qty = round(curr_aggregated_qty, qty_decimal)
      if curr_aggregated_qty > 0:
        self._price_maps[update_side][curr_aggregated_price] = \
            PriceQty(curr_aggregated_price, curr_aggregated_qty)
    else:
      raise ValueError()

  def update_aggregated_book(self, update_side: BookUpdateSide, price, qty, qty_decimal=8):
    self.request_timestamp = self._original_book.request_timestamp
    self.fetched_timestamp = self._original_book.fetched_timestamp
    self.exchange_timestamp = self._original_book.exchange_timestamp

    if update_side == BookUpdateSide.BID:
      aggregated_price = to_aggregated_price_bid(price, self._aggregated_ticksize)
    elif update_side == BookUpdateSide.ASK:
      aggregated_price = to_aggregated_price_ask(price, self._aggregated_ticksize)
    else:
      raise ValueError(update_side)
    price_qty = self._price_maps[update_side].get(aggregated_price, PriceQty(aggregated_price, 0))
    new_price_qty = PriceQty(aggregated_price, round(price_qty.qty + qty, qty_decimal))
    if new_price_qty.qty <= 0:
      if aggregated_price in self._price_maps[update_side]:
        del self._price_maps[update_side][aggregated_price]
    else:
      self._price_maps[update_side][aggregated_price] = new_price_qty
